package com.joker.memory.controller;

import com.alibaba.cloud.ai.memory.jdbc.MysqlChatMemoryRepository;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

/**
 * 使用mysql数据库存储聊天记录
 *
 * @author 郭鹏伟
 */
@RestController
@RequestMapping("/advisor/memory/mysql")
public class MysqlMemoryController {

    //AI对话客户端
    private final ChatClient chatClient;
    //储存最大会话数
    private final int MAX_MESSAGES = 100;
    // AI聊天窗口管理-管理对话上下文窗口的组件
    private final MessageWindowChatMemory messageWindowChatMemory;

    public MysqlMemoryController(ChatClient.Builder builder, MysqlChatMemoryRepository mysqlChatMemoryRepository) {
        this.messageWindowChatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(mysqlChatMemoryRepository)
                .maxMessages(MAX_MESSAGES)
                .build();
        this.chatClient = builder
                .defaultAdvisors(
                    MessageChatMemoryAdvisor.builder(messageWindowChatMemory).build()
                ).build();
    }

    @GetMapping("/call")
    public String call(
            @RequestParam(value = "message", defaultValue = "你好！我叫joker,请记得我的名字。") String message,
            @RequestParam(value = "id", defaultValue = "G007") String conversationId
    ) {
        return chatClient.prompt(message)
                .advisors(
                        a -> a
                                .param(CONVERSATION_ID, conversationId)
                                .param("CHAT_MEMORY_RETRIEVE_SIZE_KEY", MAX_MESSAGES)
                )
                .call()
                .content();
    }

    @GetMapping("/stream")
    public Flux<String> stream(
            @RequestParam(value = "message", defaultValue = "你好！我叫joker,请记得我的名字。") String message,
            @RequestParam(value = "id", defaultValue = "G007") String conversationId
    ) {
        return chatClient.prompt()
                .user(message)
                .advisors(
                        a -> a
                                .param(CONVERSATION_ID, conversationId)
                                .param("CHAT_MEMORY_RETRIEVE_SIZE_KEY", MAX_MESSAGES)
                )
                .stream()
                .content();
    }

    @GetMapping("/getMessageById")
    public List<Message> getMessageById(@RequestParam(value = "id", defaultValue = "G007") String conversationId) {
        return messageWindowChatMemory.get(conversationId);
    }

}
